ローカルのJupyter環境からAmazon Bedrockにお願いしてSQLのリファクタリングをしてもらった

みなさん、こんにちは。

クルトンです！

今回は、前回ブログと前々回ブログに書いたSQLが良い書き方なのか気になったので、ローカルマシンからAmazon Bedrockに聞いてみました。

前提

本ブログでは、2種類の実行環境を利用しています。

• Amazon Bedrockを動作させるために、poetryを使ってJupyter Lab環境を準備します。
• Amazon Bedrockに教えてもらった内容を実行するために、Google Colab(ランタイムはCPU)を使っています。

「poetryで準備したいが、poetryがローカルマシンに入っていないよ」という方は以下にナイスな記事がございますのでご参照ください。

ローカル環境準備

poetryを使って実行環境を用意していきます。

まずは実行するためのフォルダを新規作成して移動しましょう。カレントディレクトリ内に作業フォルダを作って良い場合は以下のようにすると良いです。

mkdir work && cd work

作業フォルダを作ったら、以下の順番に実行します。

1. poetry init
   • 全ての質問項目をEnter押していけばOKです
2. poetry add boto3
   • 自分の環境では1.28.60が入りました
3. poetry add jupyterlab
   • Jupyter Labを立ち上げられるようになります
4. touch app.ipynb
   • Amazon Bedrockに質問するためのノートブックファイルを準備します(GUIで実行してもOKです)

以上でAmazon Bedrockに質問するためのローカル環境を準備できました。

あとはpoetry run jupyter labを実行すれば、ノートブック形式で実行できる環境が手に入ります。

Amazon Bedrockに聞いてみた

質問していくにあたって、とても参考になったブログを貼っておきます。プロンプト以外は同一のコードを使用していますので、ご参照ください。

ダウンサンプリングに使うSQL

最初に前々回ブログで作ったダウンサンプリングするSQLについて以下のようなpromptを作って質問してみました。

prompt = """
あなたはSQLの専門家として振る舞ってください。

Human:
以下に書いたSQLの解説と、可能であるならリファクタリングをお願いします。
WITH tbl AS(
  SELECT
    *, ROW_NUMBER() OVER(PARTITION BY species ORDER BY RANDOM()) as random_sampling
  FROM
      penguins.csv
  WHERE
    species='Adelie'
  UNION ALL
  SELECT
    *, ROW_NUMBER() OVER(PARTITION BY species ORDER BY RANDOM()) as random_sampling
  FROM
      penguins.csv
  WHERE
    species='Chinstrap'
  UNION ALL
  SELECT
    *, ROW_NUMBER() OVER(PARTITION BY species ORDER BY RANDOM()) as random_sampling
  FROM
      penguins.csv
  WHERE
    species='Gentoo'
)
, all_cnt_table AS(
  SELECT
    species ,count(*) as species_count
  FROM
    penguins.csv
  GROUP BY
    species
)
, min_cnt_table AS(
  SELECT
    MIN(species_count) as min_count
  FROM
    all_cnt_table
)
SELECT
  species
  ,island
  ,bill_length_mm
  ,bill_depth_mm
  ,flipper_length_mm
  ,body_mass_g
  ,sex
FROM
  tbl, min_cnt_table
WHERE
  random_sampling <= min_count

Assistant:
"""

以下のような返答が返ってきました。

 このSQLはペンギンのデータをランダムサンプリングしています。

まずCTEのtblで、各種類のペンギンをランダムに並べ替えています。

次にall_cnt_tableで、各種類のペンギンの総数をカウントしています。

min_cnt_tableでは、種類ごとの最小の数を抽出しています。

そしてtblとmin_cnt_tableをJOINし、random_sampling列がmin_count以下の行のみを抽出しています。

これにより、各種類のペンギンから最小の数だけランダムに抽出することができます。

リファクタリングとしては以下のようにできると思います。

1. CTEの代わりにサブクエリやテンポラリテーブルを使う

2. 種類ごとに処理を分けず、WHEREでフィルタリングする

3. テーブルJOINではなく、サブクエリで最小数を参照する

例:

SELECT
  species, 
  island,
  bill_length_mm, 
  bill_depth_mm, 
  flipper_length_mm,
  body_mass_g,
  sex
FROM (
  SELECT
    *, 
    ROW_NUMBER() OVER(PARTITION BY species ORDER BY RANDOM()) AS random_sampling
  FROM penguins
) t
WHERE random_sampling <= (
  SELECT MIN(cnt) 
  FROM (
    SELECT species, COUNT(*) AS cnt
    FROM penguins
    GROUP BY species
  ) t
)

このようにリファクタリングすることで、よりシンプルかつパフォーマンスの良いSQLにできると思います。

細かいところですが、日本語の最初の始まり方として、段落を一つ下げて「このSQLはペンギンのデータをランダムサンプリングしています。」と書かれているのが良い感じですね。

肝心の返答内容の中身についてですが、どういうテーブルデータなのか全く説明がなかったにも関わらず、こちらの実装意図を汲んでくれています。

SQLについても教えてくれているので、以下のSQLをGoogle Colab上で実行してみます。

SELECT
  species, 
  island,
  bill_length_mm, 
  bill_depth_mm, 
  flipper_length_mm,
  body_mass_g,
  sex
FROM (
  SELECT
    *, 
    ROW_NUMBER() OVER(PARTITION BY species ORDER BY RANDOM()) AS random_sampling
  FROM penguins
) t
WHERE random_sampling <= (
  SELECT MIN(cnt) 
  FROM (
    SELECT species, COUNT(*) AS cnt
    FROM penguins
    GROUP BY species
  ) t
)

実行してみると、(duckdb.CatalogException) Catalog Error: Table with name penguins does not exist!というエラーが出ました。.csvが抜けているようでした。

SQL文の中にある、2つのFROM句にpenguins.csvと修正して再び実行してみると、以下の画像の通り、実行が上手くいきました。

念のため、変数に実行結果を格納して中身を確認しましたが、元のSQLと同じく204行分のデータが得られていました。

別の書き方が分かったので勉強になりました。

最頻値を求めるSQL

基本統計量の中では、自前で実装した最頻値に関するSQLについてAmazon Bedrockに聞いてみました。

promptは以下の通りです。

prompt = """
あなたはSQLの専門家として振る舞ってください。

Human:
以下に書いたSQLの解説と、可能であるならリファクタリングをお願いします。
WITH cnt_tbl AS(
  SELECT
    species
    , bill_length_mm
    , COUNT(*) AS cnt
  FROM
    penguins.csv
  WHERE
    species='Gentoo'
  GROUP BY
    bill_length_mm, species
)
, max_cnt_tbl AS(
  SELECT
    MAX(cnt) as max_cnt
  FROM
    cnt_tbl
)
SELECT
  species
  , bill_length_mm
  , max_cnt as mode
FROM
  cnt_tbl, max_cnt_tbl
WHERE
  cnt>=max_cnt

Assistant:
"""

返答は以下の内容でした。

 このSQLはGentooペンギンのくちばしの長さ(bill_length_mm)の最頻値(mode)を求めるクエリです。

まずcnt_tblではGentooペンギンのデータをくちばしの長さごとにグループ化し、カウントを取得しています。

次にmax_cnt_tblでcnt_tblのカウントの最大値を求めています。

そして最後にcnt_tblとmax_cnt_tblをJOINし、カウントが最大値以上のレコードを抽出することで、最頻値を求めています。

リファクタリング案としては以下のような方法が考えられます:

- CTE(Common Table Expression)を使わずサブクエリを利用する
- MAX OVERを使うことでmax_cnt_tblは不要になる

例:

SELECT
  species,
  bill_length_mm,
  COUNT(*) AS cnt,
  MAX(COUNT(*)) OVER() AS max_cnt
FROM penguins
WHERE species = 'Gentoo'
GROUP BY species, bill_length_mm
HAVING COUNT(*) >= MAX(COUNT(*)) OVER()

この方法ではCTEを使わず1つのクエリ内で完結させています。
またOVERを使うことでmax_cntを効率的に求めています。

どちらの書き方も同じ結果を得ることができますが、CTEの利用は可読性を高める一方、OVERを使うことでパフォーマンスの改善が期待できます。目的に応じて適切な方法を選択することが重要です。

最頻値を求めるSQLクエリであると、解説してくれています。良い感じですね。

ではGoogle ColabでSQLを実行してみます。先ほどと同じく、FROM句をpenguins.csvに修正してから実行します。

実行すると、RuntimeError: (duckdb.BinderException) Binder Error: HAVING clause cannot contain window functions!というエラーが出力されました。

自分で改めて考えて、修正しても良いのですが、せっかくですのでAmazon Bedrockにどう修正したら良いか聞いてみました。

promptは以下です。

prompt = """
あなたはSQLの専門家として振る舞ってください。

Human:
以下に書いたSQLが'Binder Error: HAVING clause cannot contain window functions!'というエラーが出て実行ができないです。
どう修正したら良いですか？

SELECT
  species,
  bill_length_mm,
  COUNT(*) AS cnt,
  MAX(COUNT(*)) OVER() AS max_cnt
FROM penguins.csv
WHERE species = 'Gentoo'
GROUP BY species, bill_length_mm
HAVING COUNT(*) >= MAX(COUNT(*)) OVER()

Assistant:
"""

返答内容は以下です。

 HAVING句にはウィンドウ関数を指定できません。
このSQLはHAVING句でウィンドウ関数MAXを使っているためエラーになっています。

HAVING句で集約関数の結果をフィルタする場合は、ウィンドウ関数ではなく通常の集約関数を使う必要があります。

以下のように書き換えると実行できるようになります。

SELECT
  species,
  bill_length_mm,
  COUNT(*) AS cnt
FROM penguins.csv
WHERE species = 'Gentoo'  
GROUP BY species, bill_length_mm
HAVING COUNT(*) >= (
  SELECT MAX(cnt) FROM (
    SELECT 
      species,
      bill_length_mm,
      COUNT(*) AS cnt
    FROM penguins.csv
    WHERE species = 'Gentoo'  
    GROUP BY species, bill_length_mm
  ) t
)

子クエリで最大のCOUNT(*)値を求め、それとHAVING句で比較するようにしています。
このように修正することで、HAVING句でウィンドウ関数を使わずに目的の結果を得ることができます。

修正案を教えてくれたので早速Google Colabで実行してみます。 次の画像の通り、元のクエリと同じ結果が得られました。

終わりに

今回はAmazon Bedrockに、コードの意図とリファクタリング案を教えてもらいました。

ブログ内で実行したように、コードの意味が分からないものに関して質問しても良い感じに教えてくれそうなので、外部サービスに渡して良いかは注意しつつも、今後も使っていきたいですね。

今回はここまで。

それでは、また！